package org.korosoft.rusalad.syncnrun.tools;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

import javax.crypto.*;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.SecretKeySpec;
import java.io.*;
import java.net.*;
import java.nio.channels.SocketChannel;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.spec.KeySpec;
import java.util.Arrays;
import java.util.Collection;

/**
 * Secure socket factory.
 *
 * @author Dmitry Korotkov
 * @since 1.0
 */
public class SecureSocketFactory implements SocketFactory {
    private static final Log log = LogFactory.getLog(SecureSocketFactory.class);
    private static final String PASSWORD = "rT&6Bsze6";
    private static final byte[] KEY_SENTINEL = "SESSION_KEY".getBytes();
    public static final String CIPHER = "RSA";
    public static final String CIPHER_SESSION = "AES";
    private final PrivateKey privateKey;
    private final Collection<PublicKey> remotePublicKeys;

    public SecureSocketFactory(PrivateKey privateKey, Collection<PublicKey> remotePublicKeys) {
        this.privateKey = privateKey;
        this.remotePublicKeys = remotePublicKeys;
    }

    @Override
    public ServerSocket createServerSocket(int port) throws IOException {
        return new ServerSocket(port) {
            @Override
            public Socket accept() throws IOException {
                final Socket socket = super.accept();
                try {

                    final ObjectOutputStream outputStream = new ObjectOutputStream(socket.getOutputStream());
                    final ObjectInputStream inputStream = new ObjectInputStream(socket.getInputStream());

                    // Receive session key encoded by our public key and server private key

                    Object incomingObject = inputStream.readObject();
                    if (!(incomingObject instanceof byte[])) {
                        throw new RuntimeException("Protocol violation");
                    }

                    // Decrypt received message with our private key to get session key encoded with client's private key
                    // This encryption is chunked

                    Cipher asymmetricIncomingCipher = Cipher.getInstance(CIPHER);
                    asymmetricIncomingCipher.init(Cipher.DECRYPT_MODE, privateKey);

                    byte[] compiledData;
                    {
                        byte[] chunkedData = (byte[]) incomingObject;
                        int chunkCount = chunkedData[0];
                        if (chunkCount < 0) {
                            throw new RuntimeException("Protocol violation");
                        }
                        int off = 1;
                        ByteArrayOutputStream baos = new ByteArrayOutputStream();
                        while (chunkCount > 0) {
                            if (off + 2 > chunkedData.length) {
                                throw new RuntimeException("Protocol violation");
                            }
                            int len = (((int) chunkedData[off]) & 255) + ((((int) chunkedData[off + 1]) & 255) << 8);
                            off += 2;
                            if (off + len > chunkedData.length) {
                                throw new RuntimeException("Protocol violation");
                            }
                            baos.write(asymmetricIncomingCipher.doFinal(chunkedData, off, len));
                            off += len;
                            chunkCount--;
                        }
                        baos.close();
                        compiledData = baos.toByteArray();
                    }
                    // Try to decrypt decrypted message with public keys of known clients

                    PublicKey identifiedClientKey = null;
                    SecretKey sessionKey = null;

                    byte[] sessionKeyBytes = null;
                    for (PublicKey clientKey : remotePublicKeys) {
                        try {
                            Cipher clientIncomingCipher = Cipher.getInstance(CIPHER);
                            clientIncomingCipher.init(Cipher.DECRYPT_MODE, clientKey);
                            byte[] fullyDecodedSessionKey = clientIncomingCipher.doFinal(compiledData);
                            byte[] decodedSentinel = Arrays.copyOf(fullyDecodedSessionKey, KEY_SENTINEL.length);
                            if (Arrays.equals(decodedSentinel, KEY_SENTINEL)) {
                                identifiedClientKey = clientKey;
                                sessionKeyBytes = Arrays.copyOfRange(fullyDecodedSessionKey, KEY_SENTINEL.length, fullyDecodedSessionKey.length);
                                sessionKey = new SecretKeySpec(sessionKeyBytes, CIPHER_SESSION);
                                break;
                            }
                        } catch (Exception ignored) {
                        }
                    }

                    if (identifiedClientKey == null || sessionKeyBytes == null || sessionKey == null) {
                        throw new RuntimeException("Client authentication failed");
                    }

                    // Send out session key encrypted with our secret key and server public key

                    Cipher asymmetricOutgoingCipher = Cipher.getInstance(CIPHER);
                    asymmetricOutgoingCipher.init(Cipher.ENCRYPT_MODE, privateKey);
                    byte[] encodedSessionKey = asymmetricOutgoingCipher.doFinal(sessionKeyBytes);

                    // Encrypt encrypted session key with server public key
                    // This encryption should be multi-chinked

                    {
                        int offset = 0;
                        int chunkCount = 0;
                        ByteArrayOutputStream baos = new ByteArrayOutputStream();
                        baos.write(0); // Chunk count placeholder
                        if (encodedSessionKey.length > 14859) {
                            throw new RuntimeException("Cannot encrypt more that 14859 bytes");
                        }
                        while (offset < encodedSessionKey.length) {
                            int len = Math.min(encodedSessionKey.length - offset, 117);
                            Cipher serverOutgoingCipher2 = Cipher.getInstance(CIPHER);
                            serverOutgoingCipher2.init(Cipher.ENCRYPT_MODE, remotePublicKeys.iterator().next());
                            byte[] twiceEncodedSessionKeyChunk = serverOutgoingCipher2.doFinal(encodedSessionKey, offset, len);
                            baos.write(twiceEncodedSessionKeyChunk.length & 255);
                            baos.write((twiceEncodedSessionKeyChunk.length >> 8) & 255);
                            baos.write(twiceEncodedSessionKeyChunk);
                            offset += len;
                            chunkCount++;
                        }
                        baos.close();
                        byte[] twiceEncodedSessionKey = baos.toByteArray();
                        twiceEncodedSessionKey[0] = (byte) chunkCount;


                        // Send twice encoded session key

                        outputStream.writeObject(twiceEncodedSessionKey);
                    }

                    log.debug("Server: Secure handshake succeeded, session key is chosen.");

                    return new SymmetricKeySocket(sessionKey, socket);
                } catch (Exception e) {
                    socket.close();
                    throw new IOException(e);
                }
            }
        };
    }

    @Override
    public Socket createClientSocket(String host, int port) throws IOException {
        final Socket inner = createClientSocketInternal(host, port);
        return new Socket() {
            public void connect(SocketAddress endpoint) throws IOException {
                inner.connect(endpoint);
            }

            public void connect(SocketAddress endpoint, int timeout) throws IOException {
                inner.connect(endpoint, timeout);
            }

            public void bind(SocketAddress bindpoint) throws IOException {
                inner.bind(bindpoint);
            }

            public InetAddress getInetAddress() {
                return inner.getInetAddress();
            }

            public InetAddress getLocalAddress() {
                return inner.getLocalAddress();
            }

            public int getPort() {
                return inner.getPort();
            }

            public int getLocalPort() {
                return inner.getLocalPort();
            }

            public SocketAddress getRemoteSocketAddress() {
                return inner.getRemoteSocketAddress();
            }

            public SocketAddress getLocalSocketAddress() {
                return inner.getLocalSocketAddress();
            }

            public SocketChannel getChannel() {
                return inner.getChannel();
            }

            public InputStream getInputStream() throws IOException {
                return new CounterInputStream(inner.getInputStream());
            }

            public OutputStream getOutputStream() throws IOException {
                return new CounterOutputStream(inner.getOutputStream());
            }

            public void setTcpNoDelay(boolean on) throws SocketException {
                inner.setTcpNoDelay(on);
            }

            public boolean getTcpNoDelay() throws SocketException {
                return inner.getTcpNoDelay();
            }

            public void setSoLinger(boolean on, int linger) throws SocketException {
                inner.setSoLinger(on, linger);
            }

            public int getSoLinger() throws SocketException {
                return inner.getSoLinger();
            }

            public void sendUrgentData(int data) throws IOException {
                inner.sendUrgentData(data);
            }

            public void setOOBInline(boolean on) throws SocketException {
                inner.setOOBInline(on);
            }

            public boolean getOOBInline() throws SocketException {
                return inner.getOOBInline();
            }

            public void setSoTimeout(int timeout) throws SocketException {
                inner.setSoTimeout(timeout);
            }

            public int getSoTimeout() throws SocketException {
                return inner.getSoTimeout();
            }

            public void setSendBufferSize(int size) throws SocketException {
                inner.setSendBufferSize(size);
            }

            public int getSendBufferSize() throws SocketException {
                return inner.getSendBufferSize();
            }

            public void setReceiveBufferSize(int size) throws SocketException {
                inner.setReceiveBufferSize(size);
            }

            public int getReceiveBufferSize() throws SocketException {
                return inner.getReceiveBufferSize();
            }

            public void setKeepAlive(boolean on) throws SocketException {
                inner.setKeepAlive(on);
            }

            public boolean getKeepAlive() throws SocketException {
                return inner.getKeepAlive();
            }

            public void setTrafficClass(int tc) throws SocketException {
                inner.setTrafficClass(tc);
            }

            public int getTrafficClass() throws SocketException {
                return inner.getTrafficClass();
            }

            public void setReuseAddress(boolean on) throws SocketException {
                inner.setReuseAddress(on);
            }

            public boolean getReuseAddress() throws SocketException {
                return inner.getReuseAddress();
            }

            public void close() throws IOException {
                inner.close();
            }

            public void shutdownInput() throws IOException {
                inner.shutdownInput();
            }

            public void shutdownOutput() throws IOException {
                inner.shutdownOutput();
            }

            @Override
            public String toString() {
                return inner.toString();
            }

            public boolean isConnected() {
                return inner.isConnected();
            }

            public boolean isBound() {
                return inner.isBound();
            }

            public boolean isClosed() {
                return inner.isClosed();
            }

            public boolean isInputShutdown() {
                return inner.isInputShutdown();
            }

            public boolean isOutputShutdown() {
                return inner.isOutputShutdown();
            }

            public void setPerformancePreferences(int connectionTime, int latency, int bandwidth) {
                inner.setPerformancePreferences(connectionTime, latency, bandwidth);
            }
        };
    }

    private Socket createClientSocketInternal(String host, int port) throws IOException {
        final Socket socket = new Socket(host, port);
        final ObjectInputStream ois = new ObjectInputStream(socket.getInputStream());
        final ObjectOutputStream oos = new ObjectOutputStream(socket.getOutputStream());

        try {

            if (remotePublicKeys.size() != 1) {
                throw new IllegalStateException("Only one server key should be specified");
            }

            // Generate session key
            SecretKeyFactory factory = SecretKeyFactory.getInstance("PBKDF2WithHmacSHA1");
            KeySpec spec = new PBEKeySpec(PASSWORD.toCharArray(), new byte[]{(byte) (System.currentTimeMillis() & 256), (byte) ((System.currentTimeMillis() >> 8) & 256)
                    , (byte) ((System.currentTimeMillis() >> 16) & 256), (byte) ((System.currentTimeMillis() >> 24) & 256)
                    , (byte) ((System.currentTimeMillis() >> 32) & 256), (byte) ((System.currentTimeMillis() >> 40) & 256)
                    , (byte) ((System.currentTimeMillis() >> 48) & 256), (byte) ((System.currentTimeMillis() >> 56) & 256)}, 1024, 256);
            SecretKey tmp = factory.generateSecret(spec);
            SecretKey sessionKey = new SecretKeySpec(Arrays.copyOf(tmp.getEncoded(), 16), CIPHER_SESSION);
            byte[] sessionKeyBytes = sessionKey.getEncoded();

            byte[] sessionKeyWithSentinel = new byte[sessionKeyBytes.length + KEY_SENTINEL.length];
            System.arraycopy(KEY_SENTINEL, 0, sessionKeyWithSentinel, 0, KEY_SENTINEL.length);
            System.arraycopy(sessionKeyBytes, 0, sessionKeyWithSentinel, KEY_SENTINEL.length, sessionKeyBytes.length);

            // Send out session key encrypted with our secret key and server public key

            Cipher asymmetricOutgoingCipher = Cipher.getInstance(CIPHER);
            asymmetricOutgoingCipher.init(Cipher.ENCRYPT_MODE, privateKey);
            byte[] encodedSessionKey = asymmetricOutgoingCipher.doFinal(sessionKeyWithSentinel);

            // Encrypt encrypted session key with server public key
            // This encryption should be multi-chinked

            {
                int offset = 0;
                int chunkCount = 0;
                ByteArrayOutputStream baos = new ByteArrayOutputStream();
                baos.write(0); // Chunk count placeholder
                if (encodedSessionKey.length > 14859) {
                    throw new RuntimeException("Cannot encrypt more that 14859 bytes");
                }
                while (offset < encodedSessionKey.length) {
                    int len = Math.min(encodedSessionKey.length - offset, 117);
                    Cipher serverOutgoingCipher2 = Cipher.getInstance(CIPHER);
                    serverOutgoingCipher2.init(Cipher.ENCRYPT_MODE, remotePublicKeys.iterator().next());
                    byte[] twiceEncodedSessionKeyChunk = serverOutgoingCipher2.doFinal(encodedSessionKey, offset, len);
                    baos.write(twiceEncodedSessionKeyChunk.length & 255);
                    baos.write((twiceEncodedSessionKeyChunk.length >> 8) & 255);
                    baos.write(twiceEncodedSessionKeyChunk);
                    offset += len;
                    chunkCount++;
                }
                baos.close();
                byte[] twiceEncodedSessionKey = baos.toByteArray();
                twiceEncodedSessionKey[0] = (byte) chunkCount;


                // Send twice encoded session key

                oos.writeObject(twiceEncodedSessionKey);
            }
            // Receive session key encoded by our public key and server private key

            Object incomingObject = ois.readObject();
            if (!(incomingObject instanceof byte[])) {
                throw new RuntimeException("Protocol violation");
            }

            // Decrypt received message with our private key to get session key encoded with client's private key
            // This encryption is chunked

            Cipher asymmetricIncomingCipher = Cipher.getInstance(CIPHER);
            asymmetricIncomingCipher.init(Cipher.DECRYPT_MODE, privateKey);

            byte[] chunkedData = (byte[]) incomingObject;
            int chunkCount = chunkedData[0];
            if (chunkCount < 0) {
                throw new RuntimeException("Protocol violation");
            }
            int off = 1;
            ByteArrayOutputStream baos = new ByteArrayOutputStream();
            while (chunkCount > 0) {
                if (off + 2 > chunkedData.length) {
                    throw new RuntimeException("Protocol violation");
                }
                int len = (((int) chunkedData[off]) & 255) + ((((int) chunkedData[off + 1]) & 255) << 8);
                off += 2;
                if (off + len > chunkedData.length) {
                    throw new RuntimeException("Protocol violation");
                }
                baos.write(asymmetricIncomingCipher.doFinal(chunkedData, off, len));
                off += len;
                chunkCount--;
            }
            baos.close();
            byte[] compiledData = baos.toByteArray();


            // Try to decrypt decrypted message with server public key

            Cipher clientIncomingCipher = Cipher.getInstance(CIPHER);
            clientIncomingCipher.init(Cipher.DECRYPT_MODE, remotePublicKeys.iterator().next());
            byte[] fullyDecodedSessionKey = clientIncomingCipher.doFinal(compiledData);
            if (!Arrays.equals(sessionKeyBytes, fullyDecodedSessionKey)) {
                throw new RuntimeException("Server key authentication failed");
            }

            log.debug("Client: Secure handshake succeeded, session key is chosen.");
            return new SymmetricKeySocket(sessionKey, socket);
        } catch (Exception e) {
            ois.close();
            oos.close();
            socket.close();
            throw new IOException(e);
        }
    }

    private static byte[] writeObjectToByteArray(Object object) {
        try {
            ByteArrayOutputStream baos = new ByteArrayOutputStream();
            ObjectOutputStream oos = new ObjectOutputStream(baos);
            oos.writeObject(object);
            oos.close();
            return baos.toByteArray();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }

    }

    public static class SymmetricKeySocket extends Socket {
        private final Socket inner;
        private final InputStream inputStream;
        private final OutputStream outputStream;

        public SymmetricKeySocket(SecretKey secretKey, Socket inner) throws IOException {
            try {
                this.inner = inner;
                final Cipher incoming = Cipher.getInstance(CIPHER_SESSION);
                incoming.init(Cipher.DECRYPT_MODE, secretKey);
                final Cipher outgoing = Cipher.getInstance(CIPHER_SESSION);
                outgoing.init(Cipher.ENCRYPT_MODE, secretKey);
                inputStream = new CipherInputStream(inner.getInputStream(), incoming);
                outputStream = new CipherOutputStream(inner.getOutputStream(), outgoing);
            } catch (NoSuchAlgorithmException e) {
                throw new IOException(e);
            } catch (NoSuchPaddingException e) {
                throw new IOException(e);
            } catch (InvalidKeyException e) {
                throw new IOException(e);
            }
        }

        @Override
        public SocketChannel getChannel() {
            throw new UnsupportedOperationException("Not supported");
        }

        @Override
        public InputStream getInputStream() throws IOException {
            if (log.isTraceEnabled()) {
                log.trace("Obtaining input stream");
            }
            return inputStream;
        }

        @Override
        public OutputStream getOutputStream() throws IOException {
            if (log.isTraceEnabled()) {
                log.trace("Obtaining output stream");
            }
            return outputStream;
        }

        @Override
        public void connect(SocketAddress endpoint) throws IOException {
            inner.connect(endpoint);
        }

        @Override
        public void connect(SocketAddress endpoint, int timeout) throws IOException {
            inner.connect(endpoint, timeout);
        }

        @Override
        public void bind(SocketAddress bindpoint) throws IOException {
            inner.bind(bindpoint);
        }

        @Override
        public InetAddress getInetAddress() {
            return inner.getInetAddress();
        }

        @Override
        public InetAddress getLocalAddress() {
            return inner.getLocalAddress();
        }

        @Override
        public int getPort() {
            return inner.getPort();
        }

        @Override
        public int getLocalPort() {
            return inner.getLocalPort();
        }

        @Override
        public SocketAddress getRemoteSocketAddress() {
            return inner.getRemoteSocketAddress();
        }

        @Override
        public SocketAddress getLocalSocketAddress() {
            return inner.getLocalSocketAddress();
        }

        @Override
        public void setTcpNoDelay(boolean on) throws SocketException {
            inner.setTcpNoDelay(on);
        }

        @Override
        public boolean getTcpNoDelay() throws SocketException {
            return inner.getTcpNoDelay();
        }

        @Override
        public void setSoLinger(boolean on, int linger) throws SocketException {
            inner.setSoLinger(on, linger);
        }

        @Override
        public int getSoLinger() throws SocketException {
            return inner.getSoLinger();
        }

        @Override
        public void sendUrgentData(int data) throws IOException {
            inner.sendUrgentData(data);
        }

        @Override
        public void setOOBInline(boolean on) throws SocketException {
            inner.setOOBInline(on);
        }

        @Override
        public boolean getOOBInline() throws SocketException {
            return inner.getOOBInline();
        }

        @Override
        public void setSoTimeout(int timeout) throws SocketException {
            inner.setSoTimeout(timeout);
        }

        @Override
        public int getSoTimeout() throws SocketException {
            return inner.getSoTimeout();
        }

        @Override
        public void setSendBufferSize(int size) throws SocketException {
            inner.setSendBufferSize(size);
        }

        @Override
        public int getSendBufferSize() throws SocketException {
            return inner.getSendBufferSize();
        }

        @Override
        public void setReceiveBufferSize(int size) throws SocketException {
            inner.setReceiveBufferSize(size);
        }

        @Override
        public int getReceiveBufferSize() throws SocketException {
            return inner.getReceiveBufferSize();
        }

        @Override
        public void setKeepAlive(boolean on) throws SocketException {
            inner.setKeepAlive(on);
        }

        @Override
        public boolean getKeepAlive() throws SocketException {
            return inner.getKeepAlive();
        }

        @Override
        public void setTrafficClass(int tc) throws SocketException {
            inner.setTrafficClass(tc);
        }

        @Override
        public int getTrafficClass() throws SocketException {
            return inner.getTrafficClass();
        }

        @Override
        public void setReuseAddress(boolean on) throws SocketException {
            inner.setReuseAddress(on);
        }

        @Override
        public boolean getReuseAddress() throws SocketException {
            return inner.getReuseAddress();
        }

        @Override
        public void close() throws IOException {
            inner.close();
        }

        @Override
        public void shutdownInput() throws IOException {
            inner.shutdownInput();
        }

        @Override
        public void shutdownOutput() throws IOException {
            inner.shutdownOutput();
        }

        @Override
        public boolean isConnected() {
            return inner.isConnected();
        }

        @Override
        public boolean isBound() {
            return inner.isBound();
        }

        @Override
        public boolean isClosed() {
            return inner.isClosed();
        }

        @Override
        public boolean isInputShutdown() {
            return inner.isInputShutdown();
        }

        @Override
        public boolean isOutputShutdown() {
            return inner.isOutputShutdown();
        }

        @Override
        public void setPerformancePreferences(int connectionTime, int latency, int bandwidth) {
            inner.setPerformancePreferences(connectionTime, latency, bandwidth);
        }
    }
}
